# coding=utf-8

from Exploit.BaseExploit import *
from concurrent.futures import ThreadPoolExecutor
from threading import Lock
from Config import config
from Exploit.DirExploit import *
from Common.common import *
from Exploit.ParseExploit import *

'''

1、遍历clear_task_list列表，循环其中的每个字典，先判断是否是self.target的目标，如果是的话继续进行下一步

2、探测规则：以ip + 遍历port的形式 进行访问探测是否存在，状态码为 broken_5里面的三种 200 206 ....

3、最后每个IP扫描完之后都作为字典进行保存到列表中，最后用write_file进行写入完成

'''


class AliveScan(Exploit):
    def __init__(self, target, clear_task_list):
        super().__init__()
        self.source = "Alive scan"
        self.headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/55.0.2883.75 Safari/537.36"}
        self.ports = {}
        self.target = target
        self.clear_task_list = clear_task_list
        self.alivescanlist = []
        self.lock = Lock()

    # 获取 title status frame
    def get_statusAndtitleAndframe(self, link):
        try:
            resp = requests.get(link, timeout=config.timeout, verify=config.verify_ssl,
                                allow_redirects=config.allow_redirects)
            detectencode = chardet.detect(resp.content)  # 利用chardet模块检测编码
            try:
                title = re.findall(r'<title>(.*?)</title>', resp.content.decode(detectencode["encoding"]), re.S)[
                    0].strip(' ').strip('\r\n').strip('\n').strip('\r')
            except:
                title = ''
            try:
                status_code = resp.status_code
            except:
                status_code = ''
            try:
                frame = resp.headers.get('X-Powered-By')
            except:
                frame = ''
            return title, status_code, frame
        except:
            pass

    # 实现拼接url + port的功能
    def gen_url_list(self, target, port):
        ports = set()

        # 这里进行取要扫描的端口类型
        if port in {'default', 'small', 'medium', 'large', 'service'}:
            ports = config.ports.get(port)

        # 生成URL
        url_list = []
        target = target.strip()

        # 这里进行端口和url的拼接，最后返回一个列表
        for port in ports:
            url = Common_url_by_port(target, int(port))
            if isinstance(url, list):
                url_list.extend(url)
            else:
                url_list.append(url)

        # 最后进行全部的返回
        return url_list

    def write_file(self, web_lists, target, page):
        workbook = openpyxl.load_workbook(abs_path + str(target) + ".xlsx")
        worksheet = workbook.worksheets[page]
        index = 0
        while index < len(web_lists):
            web = list()
            web.append(web_lists[index]['url'])
            web.append(web_lists[index]['status'])
            web.append(web_lists[index]['title'])
            web.append(web_lists[index]['frame'])
            worksheet.append(web)
            index += 1
        workbook.save(abs_path + str(target) + ".xlsx")
        workbook.close()

    def exploit(self, url_list):
        # gen_url_list之后 就需要进行requests模块进行请求来进行探测是否存活
        # print(url_list)
        for url in url_list:
            try:
                title, status_code, frame = self.get_statusAndtitleAndframe(url)
                # print(title, status_code, frame)
                # 除去不想要的状态码config.ignore_status_code
                if int(status_code) in config.ignore_status_code:
                    raise Exception

                result_dict = {
                    'url': url,
                    'title': title,
                    'status': status_code,
                    'frame': frame
                }
                # print(result_dict)
                # 进行添加操作
                self.alivescanlist.append(result_dict)

            except:
                pass

    def main(self):
        logging.info("AliveScan Start")
        '''
        测试数据:
            [
                {'subdomain': '', 'ips': '183.136.237.194', 'port': ['8443', '11680', '9443'], 'target': 'ip'},
                {'subdomain': 'mzw.zj.gov.cn', 'ips': '183.136.237.220', 'port': None, 'target': 'webdomain'}
                {'subdomain': 'yxapi.nbcc.cn', 'ips': '60.190.19.102', 'port': None, 'target': 'subdomain'}
            ]
        '''
        # print(1111111)
        p = ThreadPoolExecutor(10)
        for i in self.clear_task_list:
            # print(i)
            # 先判断要扫描的目标是否是子域名下的
            if i['target'] == 'subdomain':
                # 再去获得它当前域名解析的ip
                temp_subdomain = i['subdomain']
                temp_ip = i['ips']
                # 错误/过期的子域名就会导致解析不了ip，那么这里继续判断是否解析成功也就是域名对应的ip是否为空
                if temp_subdomain != '':
                    a_url_list = self.gen_url_list(temp_subdomain, config.url_scan_mode)  # 拼接为协议+域名+端口，默认扫描的常用的web服务的端口
                    # print(a_url_list)
                    p.submit(self.exploit, a_url_list)
                if temp_ip != '':
                    b_url_list = self.gen_url_list(temp_ip, str(config.ip_scan_mode))
                    # print(b_url_list)
                    p.submit(self.exploit, b_url_list)
        p.shutdown()
        self.lock.acquire()
        self.write_file(Common_getUniqueList(self.alivescanlist), self.target, 8)
        self.lock.release()
        DirScan(self.target, self.alivescanlist).main()
        ParseScan(self.target, self.alivescanlist).main()


if __name__ == '__main__':
    # 测试
    test_list = [
        {'subdomain': '', 'ips': '183.136.237.194', 'port': ['8443', '11680', '9443'], 'target': 'ip'},
        {'subdomain': 'mzw.zj.gov.cn', 'ips': '183.136.237.220', 'port': None, 'target': 'webdomain'},
        {'subdomain': 'yxapi.nbcc.cn', 'ips': '60.190.19.102', 'port': None, 'target': 'subdomain'}
    ]

    AliveScan('nbcc.cn', test_list).main()








